﻿using FaissOperation.FaissIndex.Native;
using FaissOperation.Interface;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;

namespace FaissOperation.FaissIndex.Index
{
    public sealed class IVFIndex : IFaissIndex
    {
        private IntPtr indexPtr = IntPtr.Zero;
        private bool disposed = false;
        private bool isTrained = false;
        private int _nlist;

        public int Dimension { get; private set; }
        public bool IsTrained
        {
            get
            {
                if (disposed) return false;

                try
                {
                    isTrained = FaissNative.faiss_Index_is_trained(indexPtr);
                    return isTrained;
                }
                catch
                {
                    return false;
                }
            }
        }

        public IVFIndex(FlatIndex quantizer, int dimension, int nlist, FaissMetricType metricType)
        {
            if (quantizer == null)
                throw new ArgumentNullException(nameof(quantizer));
            if (dimension <= 0)
                throw new ArgumentException("Dimension must be positive", nameof(dimension));
            if (nlist <= 0)
                throw new ArgumentException("nlist must be positive", nameof(nlist));

            Dimension = dimension;
            _nlist = nlist;

            var err = FaissNative.faiss_IndexIVFFlat_new(out indexPtr, quantizer.GetNativePtr(), dimension, nlist, metricType);
            if (err != 0)
                throw new Exception($"创建 IVFIndex 失败, 错误码={err}");

            if (indexPtr == IntPtr.Zero)
                throw new Exception("创建 IVFIndex 失败: 返回了空指针");
        }

        public IntPtr GetNativePtr()
        {
            ThrowIfDisposed();
            return indexPtr;
        }

        public void Train(float[][] vectors)
        {
            ThrowIfDisposed();

            if (vectors == null || vectors.Length == 0)
                throw new ArgumentException("Training vectors cannot be null or empty");

            if(vectors.Length < _nlist)
                throw new ArgumentException($"训练向量的数量 {vectors.Length} 必须大于或等于 nlist {_nlist}");

            var flat = FaissUtils.Flatten(vectors, Dimension);

            var handle = GCHandle.Alloc(flat, GCHandleType.Pinned);
            try
            {
                var result = FaissNative.faiss_Index_train(indexPtr, vectors.Length, flat);
                if (result != 0)
                    throw new Exception($"训练 IVF 索引失败, 错误码: {result}");

                isTrained = FaissNative.faiss_Index_is_trained(indexPtr);

                if (!isTrained)
                    throw new Exception("训练完成但索引仍显示为未训练状态");
            }
            finally
            {
                handle.Free();
            }
        }

        public void Add(float[] vector)
        {
            ThrowIfDisposed();

            if(!IsTrained)
                throw new InvalidOperationException("IVF 索引必须先训练才能添加数据");

            FaissUtils.SafeAdd(indexPtr, vector, Dimension);
        }

        public void BatchAdd(float[][] vectors)
        {
            ThrowIfDisposed();

            if (!IsTrained)
                throw new InvalidOperationException("IVF 索引必须先训练才能添加数据");

            FaissUtils.SafeBatchAdd(indexPtr, vectors, Dimension);
        }

        public (long[] ids, float[] distances) Search(float[] query, int topK)
        {
            ThrowIfDisposed();

            if (!IsTrained)
                throw new InvalidOperationException("IVF 索引必须先训练才能搜索");

            return FaissUtils.SafeSearch(indexPtr, query, topK, Dimension);
        }

        /// <summary>
        /// 重建索引 - 清空所有数据但保持索引结构
        /// </summary>
        public void Rebuild()
        {
            ThrowIfDisposed();
            FaissUtils.SafeRebuild(indexPtr);
        }

        private void ThrowIfDisposed()
        {
            if (disposed) throw new ObjectDisposedException(nameof(IVFIndex));
        }

        public void Dispose()
        {
            if (!disposed && indexPtr != IntPtr.Zero)
            {
                FaissNative.faiss_Index_free(indexPtr);
                indexPtr = IntPtr.Zero;
                disposed = true;
            }
            GC.SuppressFinalize(this);
        }

        ~IVFIndex() => Dispose();
    }
}
